import os
import pandas as pd
import numpy as np
from tools.select_B1.Selector import compute_bbi, compute_kdj, compute_dif

def calculate_b1_tag(df):
    """
    计算B1_tag列的值
    当某一天J值低于10时，以第二天开盘价为基准，
    记录之后第三天及以后是否日内最高先涨5%或收盘先跌5%
    """
    # 初始化B1_tag列为0
    b1_tag = [0] * len(df)
    
    # 遍历数据帧（除了最后两天，因为需要至少3天的数据进行计算）
    for i in range(len(df) - 2):
        # 检查当天J值是否低于10
        if df['j'].iloc[i] < 10:
            # 以第二天的开盘价为基准
            if i + 1 < len(df):
                base_price = df['open'].iloc[i + 1]
                # 计算5%的涨跌幅
                up_threshold = base_price * 1.05
                down_threshold = base_price * 0.97
                
                # 从第三天开始检查
                for j in range(i + 2, len(df)):
                    # 检查是否先达到上涨5%
                    if df['high'].iloc[j] >= up_threshold:
                        b1_tag[i] = 1
                        break
                    # 检查是否先达到下跌5%
                    elif df['close'].iloc[j] <= down_threshold:
                        b1_tag[i] = -1
                        break
    
    return b1_tag

def calculate_atr(df, window=14):
    """
    计算ATR (Average True Range)
    ATR反映价格波动的平均范围，不考虑方向
    """
    # 计算真实波幅（True Range）
    high_low = df['high'] - df['low']
    high_close = abs(df['high'] - df['close'].shift(1))
    low_close = abs(df['low'] - df['close'].shift(1))
    
    # 计算真实波幅（取三者最大值）
    true_range = pd.DataFrame({
        'high_low': high_low,
        'high_close': high_close,
        'low_close': low_close
    }).max(axis=1)
    
    # 计算ATR（使用滚动窗口的平均值）
    atr = true_range.rolling(window=window).mean()
    return atr

def process_tdx_data():
    """
    处理TDX数据，添加BBI和J值列，并保存到新的文件夹
    """
    # 定义输入和输出目录
    input_dir = "tdx_data"
    output_dir = "tdx_data_process"
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建目录: {output_dir}")
    
    # 获取所有CSV文件
    csv_files = [f for f in os.listdir(input_dir) if f.endswith('.csv')]
    
    print(f"找到 {len(csv_files)} 个CSV文件")
    
    # 处理每个文件
    for i, filename in enumerate(csv_files):
        try:
            # 读取数据
            file_path = os.path.join(input_dir, filename)
            df = pd.read_csv(file_path)
            
            # 确保列名是小写的
            df.columns = df.columns.str.lower()
            
            # 计算BBI值
            bbi = compute_bbi(df)
            
            # 计算KDJ值
            df_with_kdj = compute_kdj(df)
            
            # 计算MACD的DIF白线
            dif = compute_dif(df)
            
            # 添加BBI和J值到数据框，并保留两位小数
            df['bbi'] = bbi.round(2)
            df['j'] = df_with_kdj['J'].round(2)
            df['dif'] = dif.round(2)
            
            # 计算成交量与价格的交叉乘积，并保留两位小数
            df['Vol_x_Price_Change'] = (df['volume'] * (df['close'] - df['close'].shift(1))).round(2)
            
            # 计算相对成交量（当前成交量与过去20日平均成交量的比值）
            df['Relative_Volume'] = (df['volume'] / df['volume'].rolling(window=20).mean()).round(2)
            
            # 计算ATR和ATR_percentage
            atr = calculate_atr(df, window=14)
            df['ATR'] = atr.round(2)
            df['ATR_percentage'] = (atr / df['close'].shift(1) * 100).round(2) # 转化为百分比回报
            
            # 计算MA_Trend_Signal（1：上升趋势（MA20 > MA60），-1：下降趋势（MA20 < MA60））
            ma20 = df['close'].rolling(window=20).mean()
            ma60 = df['close'].rolling(window=60).mean()
            df['MA_Trend_Signal'] = np.where(ma20 > ma60, 1, np.where(ma20 < ma60, -1, 0))

            # 添加今日收盘价与昨日收盘价的比率
            df['Close_Change_Rate'] = (df['close'] / df['close'].shift(1)).round(2)
            
            # 添加今日成交量与昨日成交量的比率
            df['Volume_Change_Rate'] = (df['volume'] / df['volume'].shift(1)).round(2)
            
            # 计算B1_tag值
            df['B1_tag'] = calculate_b1_tag(df)
            
            # 保存处理后的数据
            output_path = os.path.join(output_dir, filename)
            df.to_csv(output_path, index=False)
            
            # 显示进度
            if (i + 1) % 100 == 0 or (i + 1) == len(csv_files):
                print(f"已处理 {i + 1}/{len(csv_files)} 个文件: {filename}")
                
        except Exception as e:
            print(f"处理文件 {filename} 时出错: {e}")
    
    print(f"数据处理完成，结果保存在 {output_dir} 目录中")

if __name__ == "__main__":
    process_tdx_data()